Complete binary tree inserter

Time: ctor:O(N), insert:O(1), get_root:O(1); Space: O(N); medium

A complete binary tree is a binary tree in which every level, except possibly the last, is completely filled, and all nodes are as far left as possible.

Write a data structure CBTInserter that is initialized with a complete binary tree and supports the following operations:

• CBTInserter(TreeNode root) initializes the data structure on a given tree with head node root;

• CBTInserter.insert(int v) will insert a TreeNode into the tree with value node.val = v so that the tree remains complete, and returns the value of the parent of the inserted TreeNode;

• CBTInserter.get_root() will return the head node of the tree.

Example 1:

Input/Output:

• s = CBTInserter(1) -> None

• s.insert(2) -> 1

• s.get_root -> [1,2]

Example 2:

Input/Output:

• s = CBTInserter([1,2,3,4,5,6]) -> None

• s.insert(7) -> 3

• s.insert(8) -> 4

• s.get_root() -> [1,2,3,4,5,6,7,8]

Notes:

• The initial given tree is complete and contains between 1 and 1000 nodes.

• CBTInserter.insert is called at most 10000 times per test case.

• Every value of a given or inserted node is between 0 and 5000.

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

Auxiliary Tools¶

from graphviz import Graph

class TreeTasks(object):
    def visualize_tree(self, tree):
        def add_nodes_edges(tree, dot=None):
            # Create Graph (not Digraph) object
            if dot is None:
                dot = Graph()
                dot.node(name=str(tree), label=str(tree.val))
            # Add nodes
            if tree.left:
                dot.node(name=str(tree.left), label="."+str(tree.left.val))
                dot.edge(str(tree), str(tree.left))
                dot = add_nodes_edges(tree.left, dot=dot)
            if tree.right:
                dot.node(name=str(tree.right), label=str(tree.right.val)+".")
                dot.edge(str(tree), str(tree.right))
                dot = add_nodes_edges(tree.right, dot=dot)
            return dot
        # Add nodes recursively and create a list of edges
        dot = add_nodes_edges(tree)
        # Visualize the graph
        display(dot)
        return dot

class CBTInserter(object):

    def __init__(self, root):
        """
        :type root: TreeNode
        """
        self.__tree = [root]
        for i in self.__tree:
            if i.left:
                self.__tree.append(i.left)
            if i.right:
                self.__tree.append(i.right)

    def insert(self, v):
        """
        :type v: int
        :rtype: int
        """
        n = len(self.__tree)
        self.__tree.append(TreeNode(v))
        if n % 2:
            self.__tree[(n-1)//2].left = self.__tree[-1]
        else:
            self.__tree[(n-1)//2].right = self.__tree[-1]
        return self.__tree[(n-1)//2].val

    def get_root(self):
        """
        :rtype: TreeNode
        """
        return self.__tree[0]

root = TreeNode(1)
c = CBTInserter(root)

assert c.insert(2) == 1
tree = c.get_root()
t = TreeTasks()
dot = t.visualize_tree(tree)
assert tree.val == 1
assert tree.left.val == 2

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.left = TreeNode(6)
c = CBTInserter(root)
assert c.insert(7) == 3
assert c.insert(8) == 4
tree = c.get_root()
t = TreeTasks()
dot = t.visualize_tree(tree)

See also:

https://leetcode.com/problems/complete-binary-tree-inserter

https://www.lintcode.com/problem/complete-binary-tree-inserter/description